from fastapi import FastAPI
from pydantic import BaseModel
import uvicorn
import argparse
from faiss_search import FaissSearch
from fastapi import Body, FastAPI, File, Form, Query, UploadFile, WebSocket
from fastapi.middleware.cors import CORSMiddleware
from typing import Union
from io import BytesIO

import pydantic
from utils import *

import logging



class BaseResponse(BaseModel):
    code: int = pydantic.Field(200, description="HTTP status code")
    msg: str = pydantic.Field("success", description="HTTP status message")
    size_before: int = Body(..., description="操作之前数据库大小", example=0)
    size_after: int = Body(..., description="操作之后数据库大小", example=0)

class InfoResponse(BaseModel):
    code: int = 200
    msg: str = "success"
    info: dict

class SearchRequest(BaseModel):
    query_text: str = Body(..., description="问题", example="工伤保险办理")
    top_k: int = Body(None, description="返回结果数量", example=5)
    threshold: float = Body(0.75, description="相似度阈值", example=50.0)


class SearchResponse(BaseModel):
    code: int = 200
    msg: str = "success"
    ids: list
    scores:list
    texts: list
    # similarities: list
    results: dict
    results_before_filter: int
    results_after_filter: int



async def get_info(faiss, ):
    try:
        info = faiss.get_info()
        return InfoResponse(code=200, msg="success", info=info)
    except Exception as e:
        msg = f"fail: {e}"
        return InfoResponse(code=500, msg=msg, info={})


